#include "config.h"

#include "lstatfit.h"

#define GAP	0.01

ecm_double
beta_goodness(msl_fit_t *fit)
{
	ecm_uint32	i;
	ecm_double	*yvalues = fit->yvalues;
	ecm_double	sq_diffs = 0;
	ecm_double	B, alpha, beta;

	B = gsl_vector_get(fit->s->x, 0);
	alpha = gsl_vector_get(fit->s->x, 1);
	beta = gsl_vector_get(fit->s->x, 2);

	for (i = 0; i <= 300; i++) {
		double	t = (((ecm_double)i) + GAP) / (300 + GAP * 2);
		ecm_double	prob;

		prob = 1 / B * pow(t, alpha - 1) * pow(1 - t, beta - 1);

		sq_diffs += pow(yvalues[i] - prob, 2);
	}

	return sq_diffs;
}

void
beta_dist_report(FILE *fp, msl_fit_t *fit, ecm_uint32 xdiff)
{
	ecm_uint32	i;
	ecm_double	B, alpha, beta;

	B = gsl_vector_get(fit->s->x, 0);
	alpha = gsl_vector_get(fit->s->x, 1);
	beta = gsl_vector_get(fit->s->x, 2);

	ecm_fprintf(fp, "# B: %lf alpha:%lf beta:%lf\n", B, alpha, beta);
	
	for (i = 0; i <= 300; i++) {
		double	t = (((ecm_double)i) + GAP) / (300 + GAP * 2);
		ecm_double	prob;

		prob = 1 / B * pow(t, alpha - 1) * pow(1 - t, beta - 1);

		ecm_fprintf(fp, "%lf %8.7f\n", ((ecm_double)i) / 3 + xdiff, prob);
	}
}

int
beta_dist_f(const gsl_vector *x, void *data, gsl_vector *f)
{
	msl_fit_t	*fit = (msl_fit_t *)data;
	ecm_uint32	n = fit->n;
	ecm_double	*xvalues = fit->xvalues;
	ecm_double	*yvalues = fit->yvalues;
	ecm_double	range = xvalues[fit->n - 1] - xvalues[0];
     
	double	B = gsl_vector_get(x, 0);
	double	alpha = gsl_vector_get(x, 1);
	double	beta = gsl_vector_get(x, 2);

	ecm_uint32	i;
     
	for (i = 0; i < n; i++)	{
		double	t = (xvalues[i] - xvalues[0] + GAP) / (range + GAP * 2);
		double	Yi = 1 / B * pow(t, alpha - 1) * pow(1 - t, beta - 1);
		gsl_vector_set(f, i, Yi - yvalues[i]);
	}

	return GSL_SUCCESS;
}

int
beta_dist_df(const gsl_vector *x, void *data, gsl_matrix *J)
{
	msl_fit_t	*fit = (msl_fit_t *)data;
	ecm_uint32	n = fit->n;
	ecm_double	*xvalues = fit->xvalues;
	ecm_double	range = xvalues[fit->n - 1] - xvalues[0];
	double	B = gsl_vector_get(x, 0);
	double	alpha = gsl_vector_get(x, 1);
	double	beta = gsl_vector_get(x, 2);

	ecm_uint32	i;
     
	for (i = 0; i < n; i++)	{
		double	t = (xvalues[i] - xvalues[0] + GAP) / (range + GAP * 2);
		double	dB, dalpha, dbeta;

		dB = -1 / pow(B, 2) * pow(t, alpha - 1) * pow(1 - t, beta - 1);
		dalpha = 1 / B * pow(t, alpha - 1) * log(t) * pow(1 - t, beta - 1);
		dbeta = 1 / B * pow(t, alpha - 1) * pow(1 - t, beta - 1) * log(1 - t);
		gsl_matrix_set(J, i, 0, dB);
		gsl_matrix_set(J, i, 1, dalpha);
		gsl_matrix_set(J, i, 2, dbeta);
	}
	return GSL_SUCCESS;
}
